from mcp.server.fastmcp import FastMCP, Context
import subprocess
import json
from typing import List, Optional

mcp = FastMCP("Checkpolicy MCP Server")

"""
面向软件：checkpolicy (SELinux 策略工具集)

所有 MCP tool 列表：
1. checkpolicy_create_module - 创建 SELinux 策略模块源文件
   参数：module_name (str), policy_content (str)
   返回：统一 JSON 结构

2. checkpolicy_compile_module - 编译 SELinux 策略模块
   参数：te_file (str), output_mod (str)
   返回：统一 JSON 结构

3. checkpolicy_package_module - 打包 SELinux 策略模块
   参数：mod_file (str), output_pp (str)
   返回：统一 JSON 结构

4. checkpolicy_install_module - 安装 SELinux 策略模块
   参数：pp_file (str)
   返回：统一 JSON 结构

5. checkpolicy_list_modules - 列出已安装的 SELinux 策略模块
   参数：pattern (str, optional) - 可选过滤模式
   返回：统一 JSON 结构
"""

@mcp.tool()
def checkpolicy_create_module(module_name: str, policy_content: str) -> dict:
    """
    创建 SELinux 策略模块源文件 (.te 文件)
    
    参数:
        module_name: 模块文件名（不含扩展名）
        policy_content: 策略文件内容
    
    返回 JSON 结构:
        {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    te_file = f"{module_name}.te"
    try:
        with open(te_file, 'w') as f:
            f.write(policy_content)
        
        command = f"创建文件 {te_file}"
        return {
            "success": True,
            "command": command,
            "exit_code": 0,
            "stdout": f"成功创建策略文件: {te_file}",
            "stderr": ""
        }
    except Exception as e:
        return {
            "success": False,
            "command": f"创建文件 {te_file}",
            "exit_code": 1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def checkpolicy_compile_module(te_file: str, output_mod: str) -> dict:
    """
    编译 SELinux 策略模块 (.te -> .mod)
    
    参数:
        te_file: 输入的 .te 策略文件
        output_mod: 输出的 .mod 模块文件
    
    返回 JSON 结构:
        {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    command = f"checkmodule -M -m -o {output_mod} {te_file}"
    try:
        result = subprocess.run(
            command.split(),
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def checkpolicy_package_module(mod_file: str, output_pp: str) -> dict:
    """
    打包 SELinux 策略模块 (.mod -> .pp)
    
    参数:
        mod_file: 输入的 .mod 模块文件
        output_pp: 输出的 .pp 策略包文件
    
    返回 JSON 结构:
        {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    command = f"semodule_package -o {output_pp} -m {mod_file}"
    try:
        result = subprocess.run(
            command.split(),
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def checkpolicy_install_module(pp_file: str) -> dict:
    """
    安装 SELinux 策略模块
    
    参数:
        pp_file: 要安装的 .pp 策略包文件
    
    返回 JSON 结构:
        {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    command = f"semodule -i {pp_file}"
    try:
        result = subprocess.run(
            command.split(),
            capture_output=True,
            text=True,
            timeout=30
        )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

@mcp.tool()
def checkpolicy_list_modules(pattern: Optional[str] = None) -> dict:
    """
    列出已安装的 SELinux 策略模块
    
    参数:
        pattern: 可选过滤模式（grep 模式）
    
    返回 JSON 结构:
        {
            "success": bool,
            "command": str,
            "exit_code": int,
            "stdout": str,
            "stderr": str
        }
    """
    base_command = "semodule -l"
    if pattern:
        command = f"{base_command} | grep {pattern}"
    else:
        command = base_command
    
    try:
        if pattern:
            # 使用 shell=True 来支持管道
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=30
            )
        else:
            result = subprocess.run(
                command.split(),
                capture_output=True,
                text=True,
                timeout=30
            )
        
        return {
            "success": result.returncode == 0,
            "command": command,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except Exception as e:
        return {
            "success": False,
            "command": command,
            "exit_code": -1,
            "stdout": "",
            "stderr": str(e)
        }

if __name__ == "__main__":
    mcp.run()